package com.tyrone.algorithm.uf;

/**
 * 并查集
 */
public class UF_Tree_Weighted {
    private int[] eleAndGroup;      //记录结点元素和该元素所在分组的标识
    private int[] sz;               //存储每个根结点对应的树中元素的个数
    private int count;              //记录并查集中数据的分组个数

    public UF_Tree_Weighted(int N){
        count=N;
        eleAndGroup= new int[N];
        sz=new int[N];
        for (int i = 0; i < N; i++) {
            eleAndGroup[i]=i;
        }
        for (int i = 0; i < N; i++) {
            sz[i]=1;
        }
    }
    /**
     * 获取当前并查集中的数据有多少个分组
     * @return
     */
    public int count(){
        return count;
    }

    /**
     * 判断并查集中元素p和元素q是否在同一分组中
     * @param p
     * @param q
     * @return
     */
    public boolean connected(int p,int q) {
        return find(p)==find(q);
    }

    /**
     * 元素p所在分组的标识符
     * @param p
     * @return
     */
    public int find(int p){
        while (true){
            if (p==eleAndGroup[p]){
                return p;
            }
            p=eleAndGroup[p];
        }

    }

    /**
     * 把p元素所在分组和q元素所在分组合并
     * @param p
     * @param q
     */
    public void union(int p,int q){
        int pRoot = find(p);
        int qRoot = find(q);
        if (pRoot==qRoot){
            return;
        }
        if (sz[pRoot]<sz[qRoot]){
            eleAndGroup[pRoot]=qRoot;
            sz[qRoot]+=sz[pRoot];
        }else {
            eleAndGroup[qRoot]=pRoot;
            sz[pRoot]+=sz[qRoot];
        }
        count--;

    }
}
